import os
from pathlib import Path
from pydantic import BaseModel
from typing import Dict, Optional, List, Any
from conf.config import settings
import json
from common.log import log_algo
from pydantic import SkipValidation

class PromptJson(BaseModel):
    messages: str
    description: Optional[str] = ""
    variables: Optional[List[Dict[str, Any]]] = []
    output_format: Optional[List[Dict[str, Any]]] = []

class Prompt():
    def __init__(self):
        self.prompts: Dict[str, Any] = {}
        self._load_prompts()
        
    def _load_prompts(self):
        """加载所有提示文件"""
        prompt_dir = Path(settings.PROMPT_CONFIG['PROMPT_DIR'])
        if not prompt_dir.exists():
            return
            
        for file in prompt_dir.glob('*.json'):
            try:
                prompt_name = file.stem
                with open(file, 'r', encoding='utf-8') as f:
                    json_content = PromptJson.model_validate(json.load(f))
                    message_path = json_content.messages
                    if message_path:
                        message_file = Path(prompt_dir) / message_path
                        with open(message_file, 'r', encoding='utf-8') as mf:
                            content = mf.read()
                            if content:
                                json_content.messages = content
                            else:
                                raise ValueError(f"Prompt file {message_file} is empty")
                    output_format = json_content.output_format
                    if output_format:
                        # 动态创建输出格式的BaseModel类
                        fields = {}
                        for field in output_format:
                            field_name = field["name"]
                            field_type = eval(field["type"])  # 将字符串类型转换为实际类型
                            fields[field_name] = (field_type, ...)  # 使用 ... 作为默认字段值
                        class_name = prompt_name.upper() + "_BaseModel"
                        output_basemodel = type(class_name, (BaseModel,), {
                            "__annotations__": {k: SkipValidation[v[0]] for k, v in fields.items()},
                            "__doc__": json_content.description, 
                            "model_config": {"arbitrary_types_allowed": True}
                        })
                        # 创建字段描述信息
                        output_description = {}
                        for field in output_format:
                            output_description[field["name"]] = field.get("description", "")
                        output_format = {
                            "output_basemodel": output_basemodel,
                            "output_description": output_description
                        }
                        json_content.output_format = output_format
                        self.prompts[prompt_name] = json_content
                    else:
                        raise ValueError(f"Prompt file {file} is missing 'output_format' field")
                    
            except Exception as e:
                # 记录错误但继续处理其他文件
                log_algo.error(f"加载提示文件 {file} 时出错: {str(e)}")
                continue

    def get_prompt(self, name: str) -> str:
        """获取指定名称的提示文本
        
        Args:
            name: 提示文本名称
            
        Returns:
            str: 提示文本内容
        """
        return self.prompts.get(name, '')
    
    def slot_replacement(self, prompt: str, **kwargs) -> str:
        """替换提示文本中的槽位
        
        Args:
            prompt: 提示文本
            **kwargs: 槽位参数
            
        Returns:
            str: 替换后的提示文本
        """
        for key, value in kwargs.items():
            # 处理 BaseModel 类的情况
            if isinstance(value, type) and issubclass(value, BaseModel):
                # 获取模型的 JSON schema
                schema = value.model_json_schema()
                # 转换为格式化的字符串
                value_str = json.dumps(schema["properties"], ensure_ascii=False, indent=2)
            else:
                value_str = str(value)
            prompt = prompt.replace(f'{{{{{key}}}}}', value_str)
        return prompt
    
prompt = Prompt()